from torchvision.models import vgg19
import torch.nn as nn
import torch

mid_layer_index=26
last_layer_index=35
  # +1是因为我们需要包含第last_layer_index层

import torchvision.models as models

class TruncatedVGG19(nn.Module):
    def __init__(self, original_vgg19, last_feature_layer_index):
        super(TruncatedVGG19, self).__init__()
        
        # 初始化时直接截断到指定层
        self.truncated_features = nn.Sequential(*list(original_vgg19.features)[:last_feature_layer_index+1])
        #self.classifier = original_vgg19.classifier
        
        # 特征存储列表和钩子字典
        self.features_list = []
        self.feature_hooks = {}

    def register_save_feature_hook(self, layer_index):
        """注册钩子以保存指定层的输出特征"""
        def save_feature_hook(module, input, output):
            self.features_list.append(output)
        
        # 确保之前的钩子已被移除，避免重复注册
        if layer_index in self.feature_hooks:
            self.feature_hooks[layer_index].remove()
        self.feature_hooks[layer_index] = self.truncated_features[layer_index].register_forward_hook(save_feature_hook)

    def remove_hooks(self):
        """移除所有已注册的特征保存钩子"""
        for handle in self.feature_hooks.values():
            handle.remove()
        self.feature_hooks.clear()
        self.features_list.clear()

    def get_saved_features(self):
        """获取通过钩子保存的所有特征"""
        features_copy = self.features_list.copy()
        self.features_list.clear()
        return features_copy

    def forward(self, x):
        x = self.truncated_features(x)
        # 如果需要，也可以包含classifier部分
        # x = self.classifier(x)
        return x
    
    def get_layer_features(self, input, layer_index):
        """
        获取指定层的特征
        :param input: 输入数据
        :param layer_index: 需要获取特征的层的索引
        :return: 特征
        """
        # 注册钩子
        self.register_save_feature_hook(layer_index)
        # 执行前向传播
        self(input)  # 这里直接调用self即调用了forward方法
        # 获取并返回特征
        features = self.get_saved_features()
        return features

def get_vgg_model(device,l2_idx=last_layer_index):
    vgg19 = models.vgg19(pretrained=True)  # 注意：pretrained=False仅用于示例
    truncated_model = TruncatedVGG19(vgg19, l2_idx).to(device)
    return truncated_model

def vgg_loss_fun(ref_vis, int_out, vgg_model: TruncatedVGG19):
    # 使用模型内的方法来注册hook并获取特征
    ref_vis_fea_l1 = vgg_model.get_layer_features(ref_vis, mid_layer_index)[0]
    int_out_fea_l1 = vgg_model.get_layer_features(int_out, mid_layer_index)[0]

    ref_vis_fea_l2 = vgg_model.get_layer_features(ref_vis, last_layer_index)[0]
    int_out_fea_l2 = vgg_model.get_layer_features(int_out, last_layer_index)[0]

    loss = nn.MSELoss()(ref_vis_fea_l1, int_out_fea_l1) + nn.MSELoss()(ref_vis_fea_l2, int_out_fea_l2)
    return loss

